import sys

sys.path.append("../Skip_Gram_NGE/")
from skip_gram_nge_model import SkipGramModel
from input_data import InputData
import torch.optim as optim
from tqdm import tqdm
import torch
import argumentparser as argumentparser

args = argumentparser.ArgumentParser()
# 上下文窗口c
WINDOW_SIZE = args.window_size
# mini-batch
BATCH_SIZE = args.batch_size
# 需要剔除的 低频词 的频
MIN_COUNT = args.min_count
# embedding维度
EMB_DIMENSION = args.embed_dimension
# 学习率
LR = args.learning_rate
# 负采样数
NEG_COUNT = args.neg_count

"""
    定义Word2Vec模型
"""


class Word2Vec:
    def __init__(self, input_file_name, output_file_name):
        """
        参数初始化
        :param input_file_name: 输入文件名
        :param output_file_name: 输出文件名
        """
        self.output_file_name = output_file_name
        # 加载输入数据
        self.data = InputData(input_file_name, MIN_COUNT)
        # 加载模型
        self.model = SkipGramModel(self.data.word_count, EMB_DIMENSION).cuda()
        # 设置学习率
        self.lr = LR
        # 设置学习率
        self.optimizer = optim.SGD(self.model.parameters(), lr=self.lr)

    def train(self):
        # self.model.load_state_dict(torch.load("../results/skipgram_nge.pkl"))
        print("SkipGram Training......")
        # 计算单词对
        pairs_count = self.data.evaluate_pairs_count(WINDOW_SIZE)
        print("pairs_count", pairs_count)
        # 计算batch总数
        batch_count = pairs_count / BATCH_SIZE
        print("batch_count", batch_count)
        process_bar = tqdm(range(int(5 * batch_count)))
        for i in process_bar:
            # 返回batch大小的正采样对
            pos_pairs = self.data.get_batch_pairs(BATCH_SIZE, WINDOW_SIZE)
            pos_w = [int(pair[0]) for pair in pos_pairs]
            pos_v = [int(pair[1]) for pair in pos_pairs]
            # 负采样
            neg_v = self.data.get_negative_sampling(pos_pairs, NEG_COUNT)

            pos_w = pos_w
            pos_v = pos_v
            neg_v = neg_v
            # 将梯度清零
            self.optimizer.zero_grad()
            # 前向传播得到损失
            loss = self.model.forward(pos_w, pos_v, neg_v)
            # 反向传播
            loss.backward()
            # 通过梯度下降进行梯度更新
            self.optimizer.step()

            # tqdm过程，以一个字典显示实验指标
            process_bar.set_postfix(loss=loss.data)
            # 手动更新进度条
            process_bar.update()
        # 保存模型的权重和偏置系数
        torch.save(self.model.state_dict(), "../results/skipgram_nge.pkl")
        # 保存embedding
        self.model.save_embedding(self.data.id2word_dict, self.output_file_name)


if __name__ == '__main__':
    w2v = Word2Vec(input_file_name='../data/text8.txt', output_file_name="../results/skip_gram_neg.txt")
    w2v.train()
